package com.kara.woodAgent.agent.wood.node;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.kara.woodAgent.agent.graph.Next;
import com.kara.woodAgent.agent.graph.Node;
import com.kara.woodAgent.agent.model.ModelProvider;
import com.kara.woodAgent.agent.tool.provider.ToolProvider;
import com.kara.woodAgent.agent.wood.context.WoodReactContext;
import com.kara.woodAgent.agent.wood.model.ExpectedMessage;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ResponseFormat;
import dev.langchain4j.model.chat.request.ResponseFormatType;
import dev.langchain4j.model.chat.response.ChatResponse;

import java.util.ArrayList;
import java.util.List;

/**
 * @author tzy
 * @version 1.0
 * @project wood_agent
 * @description 审阅节点
 * @date 2025/6/12 14:44:50
 */
public class ReviewNode implements Node<WoodReactContext> {


	private final ModelProvider modelProvider;

	private final ToolProvider.ToolServiceHolder toolProviderHolder;


	public ReviewNode(ModelProvider modelProvider, ToolProvider.ToolServiceHolder toolProviderHolder) {
		this.modelProvider = modelProvider;
		this.toolProviderHolder = toolProviderHolder;
	}

	private final String systemPrompt = """
			#### 任务
			根据用户输入的任务计划 和已经执行放入信息历史记录,审核即将要执行的工具消息是否符合 任务计划的预期，
			如果要执行的工具消息符合任务计划的预期，我将执行这个工具消息。
			如果不符合任务计划的预期，请你给出原因，我将重新进行规划
			#### 重点
			请你重点关注任务的执行流程是否合理,是否符合任务计划的预期.
			#### 返回结果
			请你以json 格式作为你的返回，包含以下的字段：
				- type: 要执行的工具消息是否符合任务计划的预期，符合预期就返回 agree ，不符合预期就返回 disagree
				- message: 当不符合任务计划的预期，不符合的原因
			""";


	@Override
	public Next execute(WoodReactContext context) {

		String plan = context.getPlan();//计划
		ExpectedMessage expectedMessage = context.getExpectedMessage();


		ChatLanguageModel chatModel = modelProvider.getChatModel();

		List<ChatMessage> expectedWindows = expectedMessage.getExpectedWindows();

		ArrayList<ChatMessage> chatMessages = new ArrayList<>(expectedWindows);

		ChatMessage lastMessage = chatMessages.removeLast();

		SystemMessage systemMessage = SystemMessage.from(systemPrompt);
		UserMessage userMessage = UserMessage.from(
				"##### 任务计划的预期: " + plan + "\n" +
						"##### 已经执行执行信息历史记录: " + JSON.toJSONString(chatMessages) + "\n" +
						"##### 要执行的工具消息: " + JSON.toJSONString(lastMessage) + "\n"
		);

		ChatRequest request = ChatRequest.builder()
				.responseFormat(ResponseFormat.builder().type(ResponseFormatType.JSON).build())
				.toolSpecifications(toolProviderHolder.toolSpecifications())
				.messages(List.of(systemMessage, userMessage)).build();


		ChatResponse chatResponse = chatModel.chat(request);
		AiMessage aiMessage = chatResponse.aiMessage();
		JSONObject object = JSON.parseObject(aiMessage.text());
		if ("agree".equals(object.getString("type"))) {
			expectedMessage.setDisAgreeMsg(null);
			context.setToolExecutionRequest(((AiMessage) lastMessage).toolExecutionRequests());
			return Next.NextNode("tool").step(aiMessage.toString());
		} else {
			expectedMessage.setDisAgreeMsg(object.getString("message"));
			return Next.NextNode("model").step(aiMessage.toString());
		}
	}
}
